# -*- coding: utf8 -*-

import numpy as np
from io import BytesIO
import zlib, bz2, lzma
from project.mysql.table import Table
from pymysql import Binary

class TableImageVectorNumpy( Table ):

    _table_name = "image_vector_numpy"

    def create( self ):
        cur = self.getCursor()

        # Create table
        cur.execute(
            "CREATE TABLE IF NOT EXISTS %s (" % self._table_name+
            "image_id INTEGER UNSIGNED PRIMARY KEY, "+
            "file_id INTEGER UNSIGNED NOT NULL, "+
            "mode CHAR(1) NOT NULL, "+
            "vector LONGBLOB NOT NULL) ENGINE %s" % self.getEngine()
        )
        try:
            cur.execute(
                "ALTER TABLE `%s` ADD INDEX files (file_id)" % self._table_name
            )
        except:
            pass
        cur.close()
        return


    def insert( self, file_id, image_id, numpy_array ):

        # print(self.__class__.__name__ + " :: insert vector: %r" % (vector,))

        v_str = BytesIO()
        np.save(v_str, numpy_array)
        v_str = v_str.getvalue()

        len_array = len(v_str)

        gz_str = zlib.compress(v_str, 4)
        len_gz = len(gz_str)

        bz_str = bz2.compress(v_str, 2)
        len_bz = len(bz_str)

        lz_str = lzma.compress(v_str, 2)
        len_lz = len(lz_str)

        code = 's'
        if len_array > len_gz:
            code = 'g'
            v_str = gz_str
            len_array = len_gz
        if len_array > len_bz:
            code = 'b'
            v_str = bz_str
            len_array = len_bz
        if len_array > len_lz:
            code = 'l'
            v_str = lz_str
            len_array = len_lz

        v_str = Binary(v_str)

        cur = self.getCursor()
        v_id = cur.execute(
            "INSERT INTO %s(" % self._table_name+
            "file_id, image_id, mode, vector"
            ") VALUES (%s, %s, %s, %s)", (
                file_id, image_id, code, v_str
            )
        )
        self.getConnection().commit()
        cur.close()
        return v_id


    def get( self, image_id ):
        cur = self.getCursor()
        cur.execute(
            "SELECT * FROM %s " % self._table_name+
            "WHERE image_id=%(id)s",
            {"id": image_id}
        )
        vector = cur.fetchone()
        cur.close()

        v_str = vector["vector"]

        code = vector["mode"]

        if code == 'g':
            v_str = zlib.decompress(v_str)
        if code == 'b':
            v_str = bz2.decompress(v_str)
        if code == 'l':
            v_str = lzma.decompress(v_str)

        vector["vector"] = np.load(BytesIO(v_str))

        return vector


    def remove_by_file(self, file_id):
        cur = self.getCursor()
        rc = cur.execute("DELETE FROM %s " % self._table_name +
                         "WHERE file_id=%(id)s", {"id": file_id})
        self.getConnection().commit()
        cur.close()
        return rc

    pass
